import torch

from downstream.model_builder import load_state_with_same_shape
from model import Classifier, Preprocessing, ResNetFCN, SegNet


def forgiving_state_restore(net, loaded_dict):
    """
    Handle partial loading when some tensors don't match up in size.
    Because we want to use models that were trained off a different
    number of classes.
    """
    loaded_dict = {k.replace("module.", ""): v for k, v in loaded_dict.items()}
    net_state_dict = net.state_dict()
    new_loaded_dict = {}
    for k in net_state_dict:
        new_k = k
        if (
            new_k in loaded_dict
            and net_state_dict[k].size() == loaded_dict[new_k].size()
        ):
            new_loaded_dict[k] = loaded_dict[new_k]
        else:
            print("Skipped loading parameter {}".format(k))
    net_state_dict.update(new_loaded_dict)
    net.load_state_dict(net_state_dict)
    return net


def make_model(config):
    """
    Build models according to what is in the config
    """
    if config["model_images"] == "segnet":
        model_images = SegNet(config, preprocessing=Preprocessing())
    elif config["model_images"] == "resnetfcn":
        model_images = ResNetFCN(config, preprocessing=Preprocessing())
    else:
        model_images = None

    model_classifier = Classifier(config)

    checkpoint = torch.load(
        "output/segnet_12_new/200323-1125/lightning_logs/version_0/checkpoints/imgmodel-epoch= 60-m_IoU= 0.8467.ckpt",
        map_location="cpu",
    )

    img_state = {
        k.replace("model_images.", ""): v
        for k, v in checkpoint["state_dict"].items()
        if k.startswith("model_images.")
    }
    model_images.load_state_dict(img_state, strict=True)

    classifier_state = {
        k.replace("model_classifier.", ""): v
        for k, v in checkpoint["state_dict"].items()
        if k.startswith("model_classifier.")
    }
    model_classifier.load_state_dict(classifier_state, strict=True)

    # checkpoint = torch.load(
    #     "imgmodel_12.ckpt", map_location="cpu")

    # img_state = {k.replace("model_images.", ""): v for k,
    #              v in checkpoint["state_dict"].items() if k.startswith("model_images.")}
    # model_images.load_state_dict(img_state, strict=False)

    return model_images, model_classifier
